CLI Implementation with Click #2107

djsaunde · 2024-11-27T20:30:12Z

Description

This PR implements a simple CLI using click. Note the following entrypoint:

    entry_points={
        "console_scripts": [
            "axolotl=axolotl.cli.main:main",
        ],
    },

This means, once pip installed, axolotl's CLI can be used as:

axolotl [command] [args]

It support the following commands:

axolotl preprocess
axolotl train
axolotl inference
axolotl shard
axolotl merge_sharded_fsdp_weights
axolotl merge_lora
axolotl fetch

The last command is new functionality which allows the user to download (really, sync) the contents of the examples/ and deepspeed_configs/ directories at the top level of the axolotl project. For example:

axolotl fetch examples
axolotl fetch deepspeed_configs

An optional argument --dest [directory] allows the user to specify an output location; otherwise, this defaults to examples/ and deepspeed_configs/ in the local directory.

Motivation and Context

This change was requested in this Notion ticket. This allows axolotl users that have installed the package via pip to use the existing CLI commands, and source users to use them with a slightly simpler interface axolotl [command] ....

How has this been tested?

Ad hoc testing on axolotlai/axolotl-cloud:main-latest image on a Runpod A40 instance on this feature branch. I sparsely tested the preprocess, train, and fetch commands.

# pre install
$ axolotl train examples/openllama-3b/lora.yml 
-bash: axolotl: command not found

$ pip install -e .
...

# works
$ axolotl preprocess examples/openllama-3b/lora.yml 

# works
$ axolotl train examples/openllama-3b/lora.yml 
...

# config file overrides? works with replacing "_" with "-" (click limitation)
$ axolotl train examples/openllama-3b/lora.yml --max-steps 1 --micro-batch-size 1 --val-set-size 0

# etc.

TODO:

Add pytest coverage for the CLI (?)
Test via PyPI package installation (?)

Types of changes

This PR adds a new file src/axolotl/cli/main.py which implements the Click CLI commands, and minor changes elsewhere to accommodate them.

src/axolotl/cli/__init__.py

src/axolotl/utils/data/rl.py

src/axolotl/cli/main.py

djsaunde · 2024-12-02T18:02:59Z

Question for the team: does it make sense to write tests against this code path? I noticed there weren't tests against the existing axolotl.cli module, but I can pad these out if desired. In the interest of moving fast, we can also backlog this.

winglian · 2024-12-03T01:14:17Z

Question for the team: does it make sense to write tests against this code path? I noticed there weren't tests against the existing axolotl.cli module, but I can pad these out if desired. In the interest of moving fast, we can also backlog this.

yeah, let's go ahead and take the time now to add some tests against the cli.

README.md

src/axolotl/cli/main.py

winglian · 2024-12-04T02:21:06Z

I think the test errors:

    def test_inference_all_options(
        cli_runner, default_config, mock_inference_deps
    ):  # pylint: disable=redefined-outer-name
        """Test inference with all possible options"""
        mock_cfg = MagicMock()
        mock_inference_deps["load_cfg"].return_value = mock_cfg
    
        cli_runner.invoke(
            cli,
            [
                "inference",
                str(default_config),
                "--load-in-8bit",
                "--base-model",
                "base/model/path",
                "--lora-model-dir",
                "lora/path",
                "--prompter",
                "my_prompter",
            ],
        )
    
        # Check all options were passed through
>       cli_args = mock_inference_deps["do_inference"].call_args[1]["cli_args"]
E       TypeError: 'NoneType' object is not subscriptable

is because some of the commands run accelerate as a new subprocess, so don't directly call the do_inference function in this scope.

djsaunde · 2024-12-04T22:06:18Z

Assuming tests pass here (they're passing in Runpod at least), this PR is ready for re-review. I think I got the testing mostly right; I'm testing mostly shallowly where I assert that args that are expected are in fact passed to the relevant train / do_inference / do_clil etc. methods in similarly named axolotl.cli.preprocess / axolotl.cli.train etc. modules. A small subset of tests are more end-to-end and should be useful for preventing unexpected bugs from making it to main.

As an aside, once we deprecate the old way of running axolotl commands on the command line (e.g., accelerate launch axolotl.cli.train ...), we can probably refactor out a good bit of cruft.

djsaunde · 2024-12-05T05:17:26Z

Adding context from Slack:

I think something weird is going on with the logging. I couldn't debug it after a few hours of trying so I've moved the tests to tests/e2e/patched for now. I think maybe the test_validations.py could be moved there instead as I think that's what's causing the issue with the caplog logic, but I'm not sure.

I also simplified some of the CLI tests so they're no longer running simple preprocess / train / shard etc. commands end to end for the sake of runtime. We have other test coverage that actually tests this functionality, so I think there's no need to duplicate.

README.md

outputs

src/axolotl/cli/utils.py

tests/e2e/patched/cli/conftest.py

src/axolotl/cli/main.py

outputs

src/axolotl/cli/train.py

djsaunde · 2024-12-05T17:48:24Z

I need to add a way to pass accelerate launch args as well.

winglian · 2024-12-05T20:55:19Z

need to update pytest calls in https://github.com/axolotl-ai-cloud/axolotl/blob/main/.github/workflows/tests.yml#L82 and on line 126 as well (and probably the nightly tests too)

.github/workflows/tests-nightly.yml

…tch other CLI test naming

Co-authored-by: Wing Lian <[email protected]>

* Initial CLI implementation with click package * Adding fetch command for pulling examples and deepspeed configs * Automating default options for CliArgs classes * Mimicking existing no config behavior * bugfix in choose_config * Updating fetch to sync instead of re-download * bugfix * isort fix * fixing yaml isort order * pre-commit fixes * simplifying argument parsing -- pass through kwargs to do_cli * make accelerate launch default for non-preprocess commands * fixing arg handling * testing None placeholder approach * removing hacky --use-gpu argument to preprocess command * Adding brief README documentation for CLI * remove (New) * Initial CLI pytest tests * progress on CLI pytest * adding inference CLI tests; cleanup * Refactor train CLI tests to remove various mocking * Major CLI test refator; adding remaining CLI codepath test coverage * pytest fixes * remove integration markers * parallelizing examples, deepspeed config downloads; rename test to match other CLI test naming * moving cli pytest due to isolation issues; cleanup * testing fixes; various minor improvements * fix * tests fix * Update tests/cli/conftest.py Co-authored-by: Wing Lian <[email protected]> --------- Co-authored-by: Dan Saunders <[email protected]> Co-authored-by: Wing Lian <[email protected]>

djsaunde added the enhancement New feature or request label Nov 27, 2024

djsaunde requested review from winglian and NanoCode012 November 27, 2024 20:30

djsaunde self-assigned this Nov 27, 2024

djsaunde commented Nov 27, 2024

View reviewed changes

src/axolotl/cli/__init__.py Show resolved Hide resolved

winglian reviewed Nov 29, 2024

View reviewed changes

src/axolotl/utils/data/rl.py Outdated Show resolved Hide resolved

winglian reviewed Nov 29, 2024

View reviewed changes

src/axolotl/cli/main.py Outdated Show resolved Hide resolved

djsaunde force-pushed the cli branch from 93588a3 to 07a2cb9 Compare December 2, 2024 13:45

djsaunde requested a review from winglian December 2, 2024 17:53

djsaunde force-pushed the cli branch from df92acf to d5885df Compare December 2, 2024 18:09

djsaunde force-pushed the cli branch from 66424ba to be09656 Compare December 3, 2024 01:28

djsaunde marked this pull request as draft December 3, 2024 17:37

winglian reviewed Dec 4, 2024

View reviewed changes

README.md Outdated Show resolved Hide resolved

winglian reviewed Dec 4, 2024

View reviewed changes

src/axolotl/cli/main.py Outdated Show resolved Hide resolved

djsaunde force-pushed the cli branch from 66428fd to fbea4de Compare December 4, 2024 20:23

djsaunde marked this pull request as ready for review December 4, 2024 20:23

djsaunde requested a review from winglian December 5, 2024 05:13

NanoCode012 reviewed Dec 5, 2024

View reviewed changes

winglian reviewed Dec 5, 2024

View reviewed changes

outputs Show resolved Hide resolved

winglian reviewed Dec 5, 2024

View reviewed changes

src/axolotl/cli/train.py Outdated Show resolved Hide resolved

djsaunde requested review from winglian and NanoCode012 December 5, 2024 21:32

winglian reviewed Dec 5, 2024

View reviewed changes

.github/workflows/tests-nightly.yml Show resolved Hide resolved

Dan Saunders and others added 24 commits December 5, 2024 22:07

bugfix

f0790b7

isort fix

422d91b

fixing yaml isort order

b11fba4

pre-commit fixes

0016e1f

simplifying argument parsing -- pass through kwargs to do_cli

49d6cda

make accelerate launch default for non-preprocess commands

44caf3a

fixing arg handling

276198f

testing None placeholder approach

8a39b7f

removing hacky --use-gpu argument to preprocess command

e5c2a51

Adding brief README documentation for CLI

507513e

remove (New)

fb5431b

Initial CLI pytest tests

ac05e29

progress on CLI pytest

363b5a0

adding inference CLI tests; cleanup

59abdc6

Refactor train CLI tests to remove various mocking

5d18989

Major CLI test refator; adding remaining CLI codepath test coverage

a8a7819

pytest fixes

ff0ddf1

remove integration markers

fe244cf

parallelizing examples, deepspeed config downloads; rename test to ma…

8f31101

…tch other CLI test naming

moving cli pytest due to isolation issues; cleanup

1cd2647

testing fixes; various minor improvements

d5f49a9

fix

5df0b2f

tests fix

cafc208

Update tests/cli/conftest.py

996bf34

Co-authored-by: Wing Lian <[email protected]>

djsaunde force-pushed the cli branch from d5d834f to 996bf34 Compare December 6, 2024 03:10

djsaunde merged commit fc973f4 into main Dec 6, 2024
10 checks passed

This was referenced Dec 6, 2024

remove accidentally included symlink #2131

Merged

Should we use master branch or stable version? #2144

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

CLI Implementation with Click #2107

CLI Implementation with Click #2107

djsaunde commented Nov 27, 2024 •

edited

Loading

djsaunde commented Dec 2, 2024

winglian commented Dec 3, 2024

winglian commented Dec 4, 2024

djsaunde commented Dec 4, 2024 •

edited

Loading

djsaunde commented Dec 5, 2024

djsaunde commented Dec 5, 2024

winglian commented Dec 5, 2024

CLI Implementation with Click #2107

CLI Implementation with Click #2107

Conversation

djsaunde commented Nov 27, 2024 • edited Loading

Description

Motivation and Context

How has this been tested?

Types of changes

djsaunde commented Dec 2, 2024

winglian commented Dec 3, 2024

winglian commented Dec 4, 2024

djsaunde commented Dec 4, 2024 • edited Loading

djsaunde commented Dec 5, 2024

djsaunde commented Dec 5, 2024

winglian commented Dec 5, 2024

djsaunde commented Nov 27, 2024 •

edited

Loading

djsaunde commented Dec 4, 2024 •

edited

Loading